import requests
import json
import time
import pandas as pd
import os
import re
from datetime import datetime

class Receiver:

    def __init__(self, channelid, authorization, proxy):
        self.df = pd.DataFrame(columns=['prompt', 'task_id', 'url', 'filename', 'is_downloaded', 'mj_done_msg'])
        self.awaiting_list = pd.DataFrame(columns=['prompt', 'task_id', 'status'])
        self.channelid = channelid
        self.authorization = authorization
        self.headers = {'authorization': self.authorization}
        self.proxies = {'http': 'http://' + proxy, 'https': 'http://' + proxy} if proxy else {}

    def retrieve_messages(self):
        r = requests.get(
            f'https://discord.com/api/v10/channels/{self.channelid}/messages?limit={10}',
            headers=self.headers,
            proxies=self.proxies)
        jsonn = json.loads(r.text)
        return jsonn

    def collecting_results(self):

        message_list = self.retrieve_messages()
        for message in message_list:

            if (message['author']['username'] == 'Midjourney Bot') and ('**' in message['content']):
                print(message)
                if len(message['attachments']) > 0:
                    # 已完成列表
                    if (message['attachments'][0]['filename'][-4:] == '.png') or (
                            '(Open on website for full quality)' in message['content']):
                        id = message['id']
                        task_id = re.findall("<@\d+>", message['content'])[0].replace("<@", "").replace(">", "")
                        prompt = message['content'].split('**')[1].split('--')[0].strip()
                        # 如果 prompt 中含有链接，则需要把链接删除掉
                        pic_url_list = re.findall(r'<http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+>', prompt)
                        if len(pic_url_list) > 0:
                            for pic_url in pic_url_list:
                                prompt = prompt.replace(pic_url, "").strip()
                        else:
                            pass
                        url = message['attachments'][0]['url']
                        filename = message['attachments'][0]['filename']
                        if id not in self.df.index:
                            self.df.loc[id] = [prompt, task_id, url, filename, 0, message]
                    # 进行中列表
                    else:
                        id = message['id']
                        task_id = re.findall("<@\d+>", message['content'])[0].replace("<@", "").replace(">", "")
                        prompt = message['content'].split('**')[1].split('--')[0].strip()
                        status = 'unknown status'
                        if ('(fast)' in message['content']) or ('(relaxed)' in message['content']):
                            try:
                                status = re.findall("(\w*%)", message['content'])[0]
                            except:
                                status = 'unknown status'
                        self.awaiting_list.loc[id] = [prompt, task_id, status]

                else:
                    id = message['id']
                    task_id = re.findall("<@\d+>", message['content'])[0].replace("<@", "").replace(">", "")
                    prompt = message['content'].split('**')[1].split('--')[0].strip()
                    status = 'unknown status'
                    if '(Waiting to start)' in message['content']:
                        status = 'Waiting to start'
                    self.awaiting_list.loc[id] = [prompt, task_id, status]